/* Zhao Zongyi, xinshengzzy@foxmail.com, 2020.11.09 */
#include "tofino/stateful_alu_blackbox.p4"
#include "tofino/pktgen_headers.p4"
#include <tofino/constants.p4>
#include <tofino/intrinsic_metadata.p4>

#include "includes/headers.p4"
#include "includes/parser.p4"
#include "includes/macro.p4"

field_list flow {
    ipv4.srcip;
    ipv4.dstip;
    ipv4.proto;
}

// hash function for main_table_1 and fingerprint generation
field_list_calculation hash_1 {
    input {
        flow;
    }
    algorithm: crc32;
    output_width: M1_IDX_WIDTH;
}

// hash function for main_table_2
field_list_calculation hash_2 {
    input {
        flow;
    }
    algorithm: crc32_extend;
    output_width: M2_IDX_WIDTH;
}

// hash function for main_table_3
field_list_calculation hash_3 {
    input {
        flow;
    }
    algorithm: crc32_lsb;
    output_width: M3_IDX_WIDTH;
}

// hash function for ancillary table
field_list_calculation hash_4 {
    input {
        flow;
    }
    algorithm: crc32_msb;
    output_width: A_IDX_WIDTH;
}

field_list_calculation digest_hash {
    input {
        flow;
    }
    algorithm: identity;
    output_width: DIGEST_WIDTH;
}

/*==========*/
// metadata for measurement program
header_type measurement_meta_t {
    fields {
    promotion: 1; // indicating variable for resubmit;
    stage: 4; // indicating variable for stage of back inserting;
    digest: 8; // digest for differentiating in ancillary table;
    M1_collision: 4; 
    M2_collision: 4;
    M3_collision: 4;
    A_pred: 4; // output predicate in main sub table 1
    flag1: 32;
    flag2: 32;
    flag3: 32;
    flag: 32; // subtract flag used for judging negative or not
    M1_cnt: 32;
    M2_cnt: 32;
    M3_cnt: 32;
    A_cnt: 8;
    srcip_pred: 4;
    dstip_pred: 4;
    proto_pred: 4;
    }
}

metadata measurement_meta_t m_meta;

action nop() {}


/*---------- Registers for M1 ----------*/
register M1_srcip
{
    width: 32;
    instance_count: M1_SIZE;
}
register M1_dstip
{
    width: 32;
    instance_count: M1_SIZE;
}
register M1_proto
{
    width: 8;
    instance_count: M1_SIZE;
}
register M1_cnt
{
    width: 32;
    instance_count: M1_SIZE;
}

/*---------- Registers for M2 ----------*/
register M2_srcip
{
    width: 32;
    instance_count: M2_SIZE;
}
register M2_dstip
{
    width: 32;
    instance_count: M2_SIZE;
}
register M2_proto
{
    width: 8;
    instance_count: M2_SIZE;
}
register M2_cnt
{
    width: 32;
    instance_count: M2_SIZE;
}

/*---------- Registers for M3 ----------*/
register M3_srcip
{
    width: 32;
    instance_count: M3_SIZE;
}
register M3_dstip
{
    width: 32;
    instance_count: M3_SIZE;
}
register M3_proto
{
    width: 8;
    instance_count: M3_SIZE;
}
register M3_cnt
{
    width: 32;
    instance_count: M3_SIZE;
}

/*---------- Registers for A ----------*/
register A
{
 width: 16; // register_lo is for the digest field, and register_hi is for the count field 
 instance_count: A_SIZE;
}

/*---------- Register for Resubmission ----------*/
register resmt_cnt {
 width: 32;
 instance_count: 1;
}

/*---------- calc_digest_t ----------*/
action calc_digest_ac()
{
    modify_field_with_hash_based_offset(m_meta.digest, 0, digest_hash, 256);
}

@pragma stage 0
table calc_digest_t
{
    actions {
        calc_digest_ac;
    }
 default_action: calc_digest_ac;
}

/*---------- forward_t ----------*/
action set_egr(egress_spec) {
  modify_field(ig_intr_md_for_tm.ucast_egress_port, egress_spec);
}

@pragma stage 0
table forward_t {
  reads {
  ig_intr_md.ingress_port: exact;
  }
  actions {
    set_egr; nop;
  }
}

/*---------- update_M1_srcip_t ----------*/
blackbox stateful_alu update_M1_srcip_bb
{
 reg: M1_srcip;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.srcip;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.srcip;

 output_value: predicate;
 output_dst: m_meta.srcip_pred;
}

action update_M1_srcip_ac()
{
    update_M1_srcip_bb.execute_stateful_alu_from_hash(hash_1);
}

@pragma stage 0
table update_M1_srcip_t
{
    actions {
        update_M1_srcip_ac;
    }
 default_action: update_M1_srcip_ac;
}

/*---------- rewrite_M1_srcip_t ----------*/
blackbox stateful_alu rewrite_M1_srcip_bb
{
 reg: M1_srcip;
 update_lo_1_value: ipv4.srcip;
}

action rewrite_M1_srcip_ac()
{
    rewrite_M1_srcip_bb.execute_stateful_alu_from_hash(hash_1);
}

@pragma stage 0
table rewrite_M1_srcip_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M1_srcip_ac;
  }
 default_action: nop;
}

/*---------- update_M1_dstip_t ----------*/
blackbox stateful_alu update_M1_dstip_bb
{
 reg: M1_dstip;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.dstip;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.dstip;

 output_value: predicate;
 output_dst: m_meta.dstip_pred;
}

action update_M1_dstip_ac()
{
    update_M1_dstip_bb.execute_stateful_alu_from_hash(hash_1);
}

@pragma stage 0
table update_M1_dstip_t
{
    actions {
        update_M1_dstip_ac;
    }
 default_action: update_M1_dstip_ac;
}

/*---------- rewrite_M1_dstip_t ----------*/
blackbox stateful_alu rewrite_M1_dstip_bb
{
 reg: M1_dstip;
 update_lo_1_value: ipv4.dstip;
}

action rewrite_M1_dstip_ac()
{
    rewrite_M1_dstip_bb.execute_stateful_alu_from_hash(hash_1);
}

@pragma stage 0
table rewrite_M1_dstip_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    rewrite_M1_dstip_ac;
    nop;
  }
 default_action: nop;
}

/*---------- update_M1_proto_t ----------*/
blackbox stateful_alu update_M1_proto_bb
{
 reg: M1_proto;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.proto;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.proto;

 output_value: predicate;
 output_dst: m_meta.proto_pred;
}

action update_M1_proto_ac()
{
    update_M1_proto_bb.execute_stateful_alu_from_hash(hash_1);
}

@pragma stage 0
table update_M1_proto_t
{
  actions {
    update_M1_proto_ac;
  }
 default_action: update_M1_proto_ac;
}

/*---------- rewrite_M1_proto_t ----------*/
blackbox stateful_alu rewrite_M1_proto_bb
{
 reg: M1_proto;
 update_lo_1_value: ipv4.proto;
}

action rewrite_M1_proto_ac()
{
    rewrite_M1_proto_bb.execute_stateful_alu_from_hash(hash_1);
}

@pragma stage 0
table rewrite_M1_proto_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    rewrite_M1_proto_ac;
    nop;
  }
 default_action: nop;
}

/*---------- update_M1_cnt_t ----------*/
blackbox stateful_alu increase_M1_cnt_bb
{
 reg: M1_cnt;
 update_lo_1_value: register_lo + 1;
}

action increase_M1_cnt_ac()
{
    increase_M1_cnt_bb.execute_stateful_alu_from_hash(hash_1);
    modify_field(m_meta.M1_collision, 0);
}

blackbox stateful_alu read_M1_cnt_bb
{
 reg: M1_cnt;
 output_value: register_lo;
 output_dst: m_meta.M1_cnt;
}

action read_M1_cnt_ac()
{
    read_M1_cnt_bb.execute_stateful_alu_from_hash(hash_1);
    modify_field(m_meta.M1_collision, 1);
}

@pragma stage 1
table update_M1_cnt_t
{
    reads {
    m_meta.srcip_pred: exact;
    m_meta.dstip_pred: exact;
    m_meta.proto_pred: exact;
    }
    actions {
      increase_M1_cnt_ac;
      read_M1_cnt_ac;
    }
}

/*---------- rewrite_M1_cnt_t ----------*/
blackbox stateful_alu rewrite_M1_cnt_bb
{
 reg: M1_cnt;
 update_lo_1_value: m_meta.A_cnt;
}

action rewrite_M1_cnt_ac()
{
    rewrite_M1_cnt_bb.execute_stateful_alu_from_hash(hash_1);
}

@pragma stage 1
table rewrite_M1_cnt_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    rewrite_M1_cnt_ac;
    nop;
  }
 default_action: nop;
}

/*---------- update_M2_srcip_t ----------*/
blackbox stateful_alu update_M2_srcip_bb
{
 reg: M2_srcip;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.srcip;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.srcip;

 output_value: predicate;
 output_dst: m_meta.srcip_pred;
}

action update_M2_srcip_ac()
{
    update_M2_srcip_bb.execute_stateful_alu_from_hash(hash_2);
}

@pragma stage 2
table update_M2_srcip_t
{
    actions {
        update_M2_srcip_ac;
    }
 default_action: update_M2_srcip_ac;
}

/*---------- rewrite_M2_srcip_t ----------*/
blackbox stateful_alu rewrite_M2_srcip_bb
{
 reg: M2_srcip;
 update_lo_1_value: ipv4.srcip;
}

action rewrite_M2_srcip_ac()
{
    rewrite_M2_srcip_bb.execute_stateful_alu_from_hash(hash_2);
}

@pragma stage 2
table rewrite_M2_srcip_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    rewrite_M2_srcip_ac;
    nop;
  }
 default_action: nop;
}

/*---------- update_M2_dstip_t ----------*/
blackbox stateful_alu update_M2_dstip_bb
{
 reg: M2_dstip;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.dstip;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.dstip;

 output_value: predicate;
 output_dst: m_meta.dstip_pred;
}

action update_M2_dstip_ac()
{
    update_M2_dstip_bb.execute_stateful_alu_from_hash(hash_2);
}

@pragma stage 2
table update_M2_dstip_t
{
    actions {
        update_M2_dstip_ac;
    }
 default_action: update_M2_dstip_ac;
}

/*---------- rewrite_M2_dstip_t ----------*/
blackbox stateful_alu rewrite_M2_dstip_bb
{
 reg: M2_dstip;
 update_lo_1_value: ipv4.dstip;
}

action rewrite_M2_dstip_ac()
{
    rewrite_M2_dstip_bb.execute_stateful_alu_from_hash(hash_2);
}

@pragma stage 2
table rewrite_M2_dstip_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M2_dstip_ac;
  }
 default_action: nop;
}

/*---------- update_M2_proto_t ----------*/
blackbox stateful_alu update_M2_proto_bb
{
 reg: M2_proto;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.proto;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.proto;

 output_value: predicate;
 output_dst: m_meta.proto_pred;
}

action update_M2_proto_ac()
{
    update_M2_proto_bb.execute_stateful_alu_from_hash(hash_2);
}

@pragma stage 2
table update_M2_proto_t
{
    actions {
        update_M2_proto_ac;
    }
    default_action: update_M2_proto_ac;
}

/*---------- rewrite_M2_proto_t ----------*/
blackbox stateful_alu rewrite_M2_proto_bb
{
 reg: M2_proto;
 update_lo_1_value: ipv4.proto;
}

action rewrite_M2_proto_ac()
{
    rewrite_M2_proto_bb.execute_stateful_alu_from_hash(hash_2);
}

@pragma stage 2
table rewrite_M2_proto_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M2_proto_ac;
  }
 default_action: nop;
}

/*---------- update_M2_cnt_t ----------*/
blackbox stateful_alu increase_M2_cnt_bb
{
 reg: M2_cnt;
 update_lo_1_value: register_lo + 1;
}

action increase_M2_cnt_ac()
{
    increase_M2_cnt_bb.execute_stateful_alu_from_hash(hash_2);
    modify_field(m_meta.M2_collision, 0);
}

blackbox stateful_alu read_M2_cnt_bb
{
 reg: M2_cnt;

 output_value: register_lo;
 output_dst: m_meta.M2_cnt;
}

action read_M2_cnt_ac()
{
  read_M2_cnt_bb.execute_stateful_alu_from_hash(hash_2);
  modify_field(m_meta.M2_collision, 1);
}

@pragma stage 3
table update_M2_cnt_t
{
    reads {
    m_meta.srcip_pred: exact;
    m_meta.dstip_pred: exact;
    m_meta.proto_pred: exact;
    }
    actions {
      read_M2_cnt_ac;
      increase_M2_cnt_ac;
    }
}

/*---------- rewrite_M2_cnt_t ----------*/
blackbox stateful_alu rewrite_M2_cnt_bb
{
 reg: M2_cnt;
 update_lo_1_value: m_meta.A_cnt;
}

action rewrite_M2_cnt_ac()
{
    rewrite_M2_cnt_bb.execute_stateful_alu_from_hash(hash_2);
}

@pragma stage 3
table rewrite_M2_cnt_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M2_cnt_ac;
  }
 default_action: nop;
}

/*---------- update_M3_srcip_t ----------*/
blackbox stateful_alu update_M3_srcip_bb
{
 reg: M3_srcip;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.srcip;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.srcip;

 output_value: predicate;
 output_dst: m_meta.srcip_pred;
}

action update_M3_srcip_ac()
{
    update_M3_srcip_bb.execute_stateful_alu_from_hash(hash_3);
}

@pragma stage 4
table update_M3_srcip_t
{
    actions {
        update_M3_srcip_ac;
    }
    default_action: update_M3_srcip_ac;
}

/*---------- rewrite_M3_srcip_t ----------*/
blackbox stateful_alu rewrite_M3_srcip_bb
{
 reg: M3_srcip;
 update_lo_1_value: ipv4.srcip;
}

action rewrite_M3_srcip_ac()
{
    rewrite_M3_srcip_bb.execute_stateful_alu_from_hash(hash_3);
}

@pragma stage 4
table rewrite_M3_srcip_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M3_srcip_ac;
  }
 default_action: nop;
}

/*---------- update_M3_dstip_t ----------*/
blackbox stateful_alu update_M3_dstip_bb
{
 reg: M3_dstip;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.dstip;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.dstip;

 output_value: predicate;
 output_dst: m_meta.dstip_pred;
}

action update_M3_dstip_ac()
{
    update_M3_dstip_bb.execute_stateful_alu_from_hash(hash_3);
}

@pragma stage 4
table update_M3_dstip_t
{
    actions {
      update_M3_dstip_ac;
    }
    default_action: update_M3_dstip_ac;
}

/*---------- rewrite_M3_dstip_t ----------*/
blackbox stateful_alu rewrite_M3_dstip_bb
{
 reg: M3_dstip;
 update_lo_1_value: ipv4.dstip;
}

action rewrite_M3_dstip_ac()
{
    rewrite_M3_dstip_bb.execute_stateful_alu_from_hash(hash_3);
}

@pragma stage 4
table rewrite_M3_dstip_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M3_dstip_ac;
  }
 default_action: nop;
}

/*---------- update_M3_proto_t ----------*/
blackbox stateful_alu update_M3_proto_bb
{
 reg: M3_proto;
 condition_lo: register_lo == 0;
 condition_hi: register_lo == ipv4.proto;

 update_lo_1_predicate: condition_lo or condition_hi;
 update_lo_1_value: ipv4.proto;

 output_value: predicate;
 output_dst: m_meta.proto_pred;
}

action update_M3_proto_ac()
{
    update_M3_proto_bb.execute_stateful_alu_from_hash(hash_3);
}

@pragma stage 4
table update_M3_proto_t
{
    actions {
        update_M3_proto_ac;
    }
    default_action: update_M3_proto_ac;
}

/*---------- rewrite_M3_proto_t ----------*/
blackbox stateful_alu rewrite_M3_proto_bb
{
 reg: M3_proto;
 update_lo_1_value: ipv4.proto;
}

action rewrite_M3_proto_ac()
{
    rewrite_M3_proto_bb.execute_stateful_alu_from_hash(hash_3);
}

@pragma stage 4
table rewrite_M3_proto_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M3_proto_ac;
  }
 default_action: nop;
}

/*---------- update_M3_cnt_t ----------*/
blackbox stateful_alu increase_M3_cnt_bb
{
 reg: M3_cnt;
 update_lo_1_value: register_lo + 1;
}

action increase_M3_cnt_ac()
{
    increase_M3_cnt_bb.execute_stateful_alu_from_hash(hash_3);
    modify_field(m_meta.M3_collision, 0);
}

blackbox stateful_alu read_M3_cnt_bb
{
 reg: M3_cnt;

 output_value: register_lo;
 output_dst: m_meta.M3_cnt;
}

action read_M3_cnt_ac()
{
    read_M3_cnt_bb.execute_stateful_alu_from_hash(hash_3);
    modify_field(m_meta.M3_collision, 1);
}

@pragma stage 5
table update_M3_cnt_t
{
    reads {
    m_meta.srcip_pred: exact;
    m_meta.dstip_pred: exact;
    m_meta.proto_pred: exact;
    }
    actions {
      increase_M3_cnt_ac;
      read_M3_cnt_ac;
    }
}

/*---------- rewrite_M3_cnt_t ----------*/
blackbox stateful_alu rewrite_M3_cnt_bb
{
 reg: M3_cnt;
 update_lo_1_value: m_meta.A_cnt;
}

action rewrite_M3_cnt_ac()
{
    rewrite_M3_cnt_bb.execute_stateful_alu_from_hash(hash_3);
}

@pragma stage 5
table rewrite_M3_cnt_t {
  reads {
  m_meta.stage: exact;
  }
  actions {
    nop;
    rewrite_M3_cnt_ac;
  }
 default_action: nop;
}

/*---------- compare_M_cnts_t ----------*/
action compare_M_cnts_ac() {
  subtract(m_meta.flag1, m_meta.M1_cnt, m_meta.M2_cnt);
  subtract(m_meta.flag2, m_meta.M1_cnt, m_meta.M3_cnt);
  subtract(m_meta.flag3, m_meta.M2_cnt, m_meta.M3_cnt);
}

@pragma stage 6
table compare_M_cnts_t {
  actions {
    compare_M_cnts_ac;
  }
 default_action: compare_M_cnts_ac;
}

/*---------- update_A_t ----------*/
blackbox stateful_alu update_A_bb
{
 reg: A;
 condition_lo: register_lo == m_meta.digest;

 update_lo_1_value: m_meta.digest;
  
 update_hi_1_predicate: condition_lo;
 update_hi_1_value: register_hi + 1;

 update_hi_2_predicate: not condition_lo;
 update_hi_2_value: 1;

 output_value: register_hi;
 output_dst: m_meta.A_cnt;
}

action update_A_ac()
{
    update_A_bb.execute_stateful_alu_from_hash(hash_4);
}

@pragma stage 6
table update_A_t
{
    actions {
      update_A_ac;
    }
    default_action: update_A_ac;
}

/*---------- compare_M_A_cnts_t ----------*/
action M1_cnt_subtract_A_cnt_ac()
{
    subtract(m_meta.flag, m_meta.M1_cnt, m_meta.A_cnt);
    modify_field(m_meta.stage, 1);
}

action M2_cnt_subtract_A_cnt_ac()
{
    subtract(m_meta.flag, m_meta.M2_cnt, m_meta.A_cnt);
    modify_field(m_meta.stage, 2);
}

action M3_cnt_subtract_A_cnt_ac()
{
    subtract(m_meta.flag, m_meta.M3_cnt, m_meta.A_cnt);
    modify_field(m_meta.stage, 3);
}

@pragma stage 7
table compare_M_A_cnts_t
{
    reads {
    m_meta.flag1: ternary;
    m_meta.flag2: ternary;
    m_meta.flag3: ternary;
    }
    actions {
	M1_cnt_subtract_A_cnt_ac;
	M2_cnt_subtract_A_cnt_ac;
	M3_cnt_subtract_A_cnt_ac;
    }
}

/*---------- resubmit_t ----------*/
blackbox stateful_alu resmt_cnt_bb
{
 reg: resmt_cnt;
 update_lo_1_value: register_lo + 1;
}

field_list resubmit_fields
{
    m_meta.promotion; //1 bits
    m_meta.stage; //4 bits
    m_meta.A_cnt; //8 bits
}

action resubmit_ac()
{
    modify_field(m_meta.promotion, 1);
    resubmit(resubmit_fields);
    resmt_cnt_bb.execute_stateful_alu();
}

@pragma stage 8
table resubmit_t
{
    reads {
        m_meta.flag: ternary;
    }
    actions {
        resubmit_ac;//highest bit is 1;
    }
}

/*==========*/
control ingress {
  // stage 0
  apply(forward_t);
  apply(calc_digest_t);
  if (0 == m_meta.promotion) {
    // stage 0
    apply(update_M1_srcip_t);
    apply(update_M1_dstip_t);
    apply(update_M1_proto_t);
    // stage 1
    apply(update_M1_cnt_t);
    if (1 == m_meta.M1_collision) {
      // stage 2
      apply(update_M2_srcip_t);
      apply(update_M2_dstip_t);
      apply(update_M2_proto_t);
      // stage 3
      apply(update_M2_cnt_t);
      if (1 == m_meta.M2_collision) {
	// stage 4
	apply(update_M3_srcip_t);
	apply(update_M3_dstip_t);
	apply(update_M3_proto_t);
	// stage 5
	apply(update_M3_cnt_t);
	if(1 == m_meta.M3_collision) {
	  // stage 6
	  apply(compare_M_cnts_t);
	  apply(update_A_t);
	  // stage 7	  
	  apply(compare_M_A_cnts_t);
	  // stage 8
	  apply(resubmit_t);
	}
      }
    }
  }
  else {
    // stage 0
    apply(rewrite_M1_srcip_t);
    apply(rewrite_M1_dstip_t);
    apply(rewrite_M1_proto_t);
    // stage 1
    apply(rewrite_M1_cnt_t);
    // stage 2
    apply(rewrite_M2_srcip_t);
    apply(rewrite_M2_dstip_t);
    apply(rewrite_M2_proto_t);
    // stage 3
    apply(rewrite_M2_cnt_t);
    // stage 4
    apply(rewrite_M3_srcip_t);
    apply(rewrite_M3_dstip_t);
    apply(rewrite_M3_proto_t);
    // stage 5
    apply(rewrite_M3_cnt_t);
  }
}

control egress
{
    // no operation
}
